#pragma once

#include <Eigen/Dense>
#include <vector>
#include <unordered_map>

namespace KalmanJoint{

//传感器编号,每个种类仅能添加一个
enum SensorNumber{
    IMU     = 0,
    RGBD_Apriltag = 1,
    RGBD_YOLO_RFRNN = 2,
};

#define QUEUE_LENGTH 500

struct HistoryData {
    Eigen::VectorXd angle;
    Eigen::VectorXd velocity;
    long double timestamp;
    std::unordered_map<SensorNumber, std::pair<bool, Eigen::VectorXd>> sensor_angle;
};

class MultiSensorKalmanFilter {
public:
    // 传感器参数结构体
    struct SensorParams {
        SensorNumber sensor_num;
        Eigen::MatrixXd sensor_H;
        Eigen::MatrixXd sensor_R;
        Eigen::VectorXd measure_date; // 每个传感器的测量值
        Eigen::VectorXd newest_measure_date; // 每个传感器的测量值
        Eigen::Matrix<double, 3, 1> sensor_residual; // 残差向量
        long double time_stamp; // 传感器时间戳,单位ms
        long double newest_time_stamp;
        bool update_flag; // 传感器数据更新标志位
        Eigen::MatrixXd sensor_B; // 误差协方差矩阵
        Eigen::VectorXd sensor_w; // 误差协方差矩阵
        bool usable_flag; // 传感器数据可用标志位
    };

    // 构造函数
    MultiSensorKalmanFilter(int state_size, int input_size, const std::vector<SensorParams>& sensors_params)
        : state_size_(state_size),
          A_(Eigen::MatrixXd::Identity(state_size, state_size)),
          B_(Eigen::MatrixXd::Identity(state_size, input_size)),
          Q_(Eigen::MatrixXd::Identity(state_size, state_size)),
          P_(Eigen::MatrixXd::Identity(state_size, state_size)),
          state_estimate_(Eigen::VectorXd::Zero(state_size)),
          input_(Eigen::VectorXd::Zero(input_size)),
          w_(Eigen::VectorXd::Ones(state_size)){
        for (const auto& params : sensors_params) {
            addSensor(params);
        }
    }

    // 创建传感器结构体
    static SensorParams CreateSensorParams(const SensorNumber& sensor_num,
                                                                                            const Eigen::MatrixXd& initial_H, const Eigen::MatrixXd& initial_R) {
        SensorParams params;
        params.sensor_H = initial_H;
        params.sensor_R = initial_R;
        params.sensor_num = sensor_num;
        return params;
    }

    // 设置状态转移矩阵
    void setTransitionMatrix(const Eigen::MatrixXd& A) {A_ = A;}

    // 设置状态转移矩阵
    void setInputMatrix(const Eigen::MatrixXd& B) {B_ = B;}

    // 设置过程噪声协方差矩阵
    void setProcessNoiseCov(const Eigen::MatrixXd& Q) {Q_ = Q;}

    Eigen::MatrixXd& getNoiseCov() {return P_;}

    // 设置初始状态估计
    void setStateEstimate(const Eigen::VectorXd& state_estimate) {
        state_estimate_ = state_estimate;
    }

    void setFilterInit(bool flag){init_ = flag;}

    void setFilterInit(bool flag, const Eigen::VectorXd& curr_estimate, const SensorNumber& base_sensor){
        init_ = flag;
        setStateEstimate(curr_estimate);
        history_data_.clear();
        for (const auto& sensor : sensor_info_) {
            const SensorNumber sensor_num = sensor.first;
            if (sensor_num != base_sensor){
                setUpdateFlag(sensor_num, false);
                setUsableFlag(sensor_num, false);
            }
        }
    }

    bool getFilterInit(){return init_;}

    // 设置输入向量
    void setInput(const Eigen::VectorXd& input) {input_ = input;}

    // 设置自适应因子
    void setAlpha(const double alpha) {alpha_ = alpha;}

    // 设置传感器最新时间戳
    void setNewestTimeStamp(const SensorNumber& sensor_num, const long double time) {
        sensor_info_[sensor_num].newest_time_stamp = time;
    }

    void setSensorTimeStamp(const SensorNumber& sensor_num, const long double time) {
        sensor_info_[sensor_num].time_stamp = time;
    }

    // 设置传感器数据标志位
    void setUpdateFlag(const SensorNumber& sensor_num, const bool flag){
        sensor_info_[sensor_num].update_flag = flag;
    }

    void setUsableFlag(const SensorNumber& sensor_num, const bool flag){
        sensor_info_[sensor_num].usable_flag = flag;
    }

    // 设置观测数据
    void setNewestData(const SensorNumber& sensor_num, const Eigen::VectorXd& m) {
        sensor_info_[sensor_num].newest_measure_date = m;
    }

    void setSensorData(const SensorNumber& sensor_num, const Eigen::VectorXd& m) {
        sensor_info_[sensor_num].measure_date = m;
    }

    Eigen::VectorXd getNewestData(const SensorNumber& sensor_num) {
        return sensor_info_[sensor_num].newest_measure_date;
    }

    Eigen::VectorXd getSensorData(const SensorNumber& sensor_num) {
        return sensor_info_[sensor_num].measure_date;
    }

    // 设置滤波器时间步长
    void setTimeStep(const long double start_time, const long double end_time) {
        time_step_ = fabs(end_time - start_time) / 1000.;
    }

    void setTimeStep(const long double dur_time) {
        time_step_ = dur_time;
    }

    long double getTimeStep() {return time_step_;}

    // KF更新函数
    void update(const SensorNumber& sensor_num);

    // AKF更新函数
    void updateAkf(const SensorNumber& sensor_num, const Eigen::VectorXd& measurement);

    void updateAllSensor();

    void updateAllData(SensorNumber base_sensor, const Eigen::VectorXd& base_sensor_data, const long double curr_time);

    // 获取当前状态估计
    Eigen::VectorXd getStateEstimate() const {return state_estimate_;}

    // 切换传感器数据标志位
    bool getUpdateFlag(const SensorNumber& sensor_num){
        return sensor_info_[sensor_num].update_flag;
    }

    bool getUsableFlag(const SensorNumber& sensor_num){
        return sensor_info_[sensor_num].usable_flag;
    }

    long double getSensorTimeStamp(const SensorNumber& sensor_num){
        return sensor_info_[sensor_num].time_stamp;
    }

    long double getNewestTimeStamp(const SensorNumber& sensor_num){
        return sensor_info_[sensor_num].newest_time_stamp;
    }

    void setResidual(const SensorNumber& sensor_num, Eigen::Matrix<double, 3, 1>& residual){
        sensor_info_[sensor_num].sensor_residual = residual;
    }

    Eigen::Matrix<double, 3, 1> getResidual(const SensorNumber& sensor_num){
        return sensor_info_[sensor_num].sensor_residual;
    }

    void addHisData(const long double curr_time);

    void updateFromHis(const SensorNumber sensor_num);

private:

    int state_size_; // 状态大小
    bool init_=false; //初始化标志位
    long double time_step_=0.01; // 滤波器时间步长,单位s
    Eigen::MatrixXd A_; // 状态转移矩阵
    Eigen::MatrixXd B_; // 输入转移矩阵
    Eigen::VectorXd input_; // 输入向量
    Eigen::MatrixXd Q_; // 过程噪声协方差矩阵
    Eigen::VectorXd w_; // 过程噪声向量
    Eigen::MatrixXd P_; // 误差协方差矩阵
    Eigen::VectorXd state_estimate_; // 状态估计
    double alpha_; // AKF自适应因子
    std::unordered_map<SensorNumber, SensorParams> sensor_info_;
    std::vector<HistoryData> history_data_;

    // 添加传感器
    void addSensor(const SensorParams& params);

    void setSensorFilterState(const SensorNumber& sensor_num){
        B_ = sensor_info_[sensor_num].sensor_B;
        w_ = sensor_info_[sensor_num].sensor_w;
    }

};
}

